# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
# --------------------------------------------------------
# References:
# DeiT: https://github.com/facebookresearch/deit
# BEiT: https://github.com/microsoft/unilm/tree/master/beit
# --------------------------------------------------------
import argparse
import datetime
import json
import numpy as np
import os
import time
from pathlib import Path

import torch
import torch.backends.cudnn as cudnn
from torch.utils.tensorboard import SummaryWriter
import torchvision.transforms as transforms
import torchvision.datasets as datasets

import timm

assert timm.__version__ == "0.3.2"  # version check
import timm.optim.optim_factory as optim_factory

import util.misc as misc
from util.misc import NativeScalerWithGradNormCount as NativeScaler

import models_mae2

from engine_pretrain import train_one_epoch
from util.MultiEpochsDataLoader import MultiEpochsDataLoader

os.environ["TF_ENABLE_ONEDNN_OPTS"] = "0"

# --------------
# Adapter
# --------------

def get_args_parser():
    parser = argparse.ArgumentParser('MAE pre-training', add_help=False)
    parser.add_argument('--batch_size', default=64, type=int,
                        help='Batch size per GPU (effective batch size is batch_size * accum_iter * # gpus')
    parser.add_argument('--epochs', default=400, type=int)
    parser.add_argument('--accum_iter', default=1, type=int,
                        help='Accumulate gradient iterations (for increasing the effective batch size under memory constraints)')

    # Model parameters
    parser.add_argument('--model', default='mae_vit_base_patch16', type=str, metavar='MODEL',
                        help='Name of model to train')

    parser.add_argument('--input_size', default=224, type=int,
                        help='images input size')

    parser.add_argument('--mask_ratio', default=0.75, type=float,
                        help='Masking ratio (percentage of removed patches).')

    parser.add_argument('--norm_pix_loss', action='store_true',
                        help='Use (per-patch) normalized pixels as targets for computing loss')
    parser.set_defaults(norm_pix_loss=False)

    # Optimizer parameters
    parser.add_argument('--weight_decay', type=float, default=0.05,
                        help='weight decay (default: 0.05)')

    parser.add_argument('--lr', type=float, default=None, metavar='LR',
                        help='learning rate (absolute lr)')
    parser.add_argument('--blr', type=float, default=1e-3, metavar='LR',
                        help='base learning rate: absolute_lr = base_lr * total_batch_size / 256')
    parser.add_argument('--min_lr', type=float, default=0., metavar='LR',
                        help='lower lr bound for cyclic schedulers that hit 0')

    parser.add_argument('--warmup_epochs', type=int, default=40, metavar='N',
                        help='epochs to warmup LR')

    # Dataset parameters
    parser.add_argument('--data_path', default=r'F:\xiaoxian(CY)\1.all_data\1.用到的数据集\3.FFHHQ数据集\yolo_crop35785\crop\xiaoxian', type=str,
                        help='dataset path')

    parser.add_argument('--output_dir', default='./output_dir',
                        help='path where to save, empty for no saving')
    parser.add_argument('--log_dir', default='./output_dir',
                        help='path where to tensorboard log')
    parser.add_argument('--device', default='cuda',
                        help='device to use for training / testing')
    parser.add_argument('--seed', default=0, type=int)
    parser.add_argument('--resume', default='',
                        help='resume from checkpoint')

    parser.add_argument('--start_epoch', default=0, type=int, metavar='N',
                        help='start epoch')
    parser.add_argument('--num_workers', default=10, type=int)
    parser.add_argument('--pin_mem', action='store_true',
                        help='Pin CPU memory in DataLoader for more efficient (sometimes) transfer to GPU.')
    parser.add_argument('--no_pin_mem', action='store_false', dest='pin_mem')
    parser.set_defaults(pin_mem=True)

    # distributed training parameters
    parser.add_argument('--world_size', default=1, type=int,
                        help='number of distributed processes')
    parser.add_argument('--local_rank', default=-1, type=int)
    parser.add_argument('--dist_on_itp', action='store_true')
    parser.add_argument('--dist_url', default='env://',
                        help='url used to set up distributed training')

    return parser


def main(args):
    # 初始化分布式训练模式，设置相关的参数。
    misc.init_distributed_mode(args)

    # 输出当前脚本所在目录和参数信息，帮助调试
    print('job dir: {}'.format(os.path.dirname(os.path.realpath(__file__))))
    print("{}".format(args).replace(', ', ',\n'))

    # 设置设备（GPU/CPU）
    device = torch.device(args.device)

    # 为了重现性，固定随机种子
    seed = args.seed + misc.get_rank()  # 加上rank，确保分布式环境下每个进程有不同的种子
    torch.manual_seed(seed)  # 设置PyTorch的随机种子
    np.random.seed(seed)  # 设置NumPy的随机种子

    # 设置cudnn.benchmark为True来优化卷积操作
    cudnn.benchmark = True

    # 定义训练数据的预处理（数据增强）方法
    transform_train = transforms.Compose([
        transforms.RandomResizedCrop(args.input_size, scale=(0.2, 1.0), interpolation=3),  # 随机裁剪并调整大小，使用bicubic插值
        transforms.RandomHorizontalFlip(),  # 随机水平翻转
        transforms.ToTensor(),  # 转为Tensor格式
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # 归一化
    ])

    # 加载训练数据集
    dataset_train = datasets.ImageFolder(args.data_path, transform=transform_train)
    print(dataset_train)

    # 如果是分布式训练，使用分布式采样器
    if True:  # args.distributed:
        num_tasks = misc.get_world_size()  # 获取全局GPU数量
        global_rank = misc.get_rank()  # 获取当前进程的rank
        sampler_train = torch.utils.data.DistributedSampler(
            dataset_train, num_replicas=num_tasks, rank=global_rank, shuffle=True
        )
        print("Sampler_train = %s" % str(sampler_train))
    else:
        # 否则，使用普通的随机采样器
        sampler_train = torch.utils.data.RandomSampler(dataset_train)

    # 如果是主进程并且设置了log目录，则创建日志文件夹，并初始化日志写入器
    if global_rank == 0 and args.log_dir is not None:
        os.makedirs(args.log_dir, exist_ok=True)
        log_writer = SummaryWriter(log_dir=args.log_dir)
    else:
        log_writer = None

    # 使用多进程加载数据
    data_loader_train = MultiEpochsDataLoader(
        dataset_train, sampler=sampler_train,
        batch_size=args.batch_size,
        num_workers=args.num_workers,
        pin_memory=args.pin_mem,
        drop_last=True,
    )

    # 定义模型
    model = models_mae2.__dict__[args.model](norm_pix_loss=args.norm_pix_loss)

    # 将模型移到目标设备上
    model.to(device)

    model_without_ddp = model  # 保存没有分布式封装的模型（在分布式训练时可能会封装）
    print("Model = %s" % str(model_without_ddp))

    # 计算有效的批次大小（考虑累积梯度和GPU数量）
    eff_batch_size = args.batch_size * args.accum_iter * misc.get_world_size()

    # 如果没有设置初始学习率，则根据基础学习率和有效批次大小计算
    if args.lr is None:  # 只指定了基础学习率
        args.lr = args.blr * eff_batch_size / 256

    # 输出相关的学习率、批次大小等信息
    print("base lr: %.2e" % (args.lr * 256 / eff_batch_size))
    print("actual lr: %.2e" % args.lr)

    print("accumulate grad iterations: %d" % args.accum_iter)
    print("effective batch size: %d" % eff_batch_size)

    # 定义优化器，使用AdamW优化器（权重衰减为0）
    optimizer = torch.optim.Adam(model_without_ddp.parameters(), lr=args.blr, weight_decay=0)

    print(optimizer)
    loss_scaler = NativeScaler()  # 定义损失缩放器（用于混合精度训练）

    # 加载已经保存的模型参数（如果存在）
    misc.load_model(args=args, model_without_ddp=model_without_ddp, optimizer=optimizer, loss_scaler=loss_scaler)

    # 输出训练开始信息
    print(f"Start training for {args.epochs} epochs")
    start_time = time.time()  # 记录训练开始时间

    # 进入训练循环，遍历每个epoch
    for epoch in range(args.start_epoch, args.epochs):
        if args.distributed:
            data_loader_train.sampler.set_epoch(epoch)  # 设置分布式训练的epoch

        # 训练一个epoch，并获取训练统计信息
        train_stats = train_one_epoch(
            model, data_loader_train,
            optimizer, device, epoch, loss_scaler,
            log_writer=log_writer,
            args=args
        )

        # 每隔20个epoch保存一次模型
        if args.output_dir and (epoch % 20 == 0 or epoch + 1 == args.epochs):
            misc.save_model(
                args=args, model=model, model_without_ddp=model_without_ddp, optimizer=optimizer,
                loss_scaler=loss_scaler, epoch=epoch)

        # 将训练结果和epoch信息保存到log文件中
        log_stats = {**{f'train_{k}': v for k, v in train_stats.items()},
                     'epoch': epoch, }

        if args.output_dir and misc.is_main_process():  # 只在主进程保存日志
            if log_writer is not None:
                log_writer.flush()  # 刷新日志写入器
            with open(os.path.join(args.output_dir, "log.txt"), mode="a", encoding="utf-8") as f:
                f.write(json.dumps(log_stats) + "\n")

    # 训练结束，计算并输出总训练时间
    total_time = time.time() - start_time
    total_time_str = str(datetime.timedelta(seconds=int(total_time)))
    print('Training time {}'.format(total_time_str))



if __name__ == '__main__':
    args = get_args_parser()
    args = args.parse_args()
    if args.output_dir:
        Path(args.output_dir).mkdir(parents=True, exist_ok=True)
    main(args)
